from collections import defaultdict
import numpy as np


class L46Inference:

    def __init__(self, corticalColumn, params):
        self.corticalColumn = corticalColumn
        self.noiseFactor = params['noiseFactor']
        self.moduleNoiseFactor = params['moduleNoiseFactor']
        self.maxSettlingTime = params['maxSettlingTime']
        self.maxTraversals = params['maxTraversals']
        self.nextMonitorToken = params['nextMonitorToken']
        self.randomLocation = params['randomLocation']
        self.useNoise = params['useNoise']
        self.noisyTrainingTime = params['noisyTrainingTime']
        self.waitForSettle = params['waitForSettle']
        self.numSensations = params['numSensations']
        # Temporal States
        self.locationRepresentations = defaultdict(list)
        self.inputRepresentations = {}
        self.learnedObjects = []
        self.locationOnObject = None
        self.monitors = {}
        self.representationSet = set()
        self.l2ObjectRepresentations = defaultdict(set)

    def learnObject(self, objectDescription):
        self.setLearnMode(True)
        self.reset()
        self.corticalColumn.activateRandomLocation()
        locationsAreUnique = True
        for iFeature, feature in enumerate(objectDescription["features"]):
            print('learning feature:', iFeature)
            self.move(feature)
            featureSDR = feature["featureValue"]
            self.sense(featureSDR)
            locationRepresentation = self.corticalColumn.getSensoryAssociatedLocationRepresentation()
            locationTuple = tuple(locationRepresentation)
            locationsAreUnique = (locationsAreUnique and locationTuple not in self.representationSet)
            if not locationsAreUnique:
                print('locations are not unique')
            self.representationSet.add(tuple(locationRepresentation))
        return locationsAreUnique

    def inferObjectSimple(self, objectDescription):
        self.reset()
        self.setLearnMode(False)
        infered = False
        currStep = 0
        for iFeature, feature in enumerate(objectDescription["features"]):
            currStep += 1
            self.move(feature)
            featureSDR = feature["featureValue"]
            self.sense(featureSDR)
            if not infered:
                currl2ObjectRep = set(self.corticalColumn.L2.activeCells)
                for key in self.l2ObjectRepresentations.keys():
                    if currl2ObjectRep <= self.l2ObjectRepresentations[key]:
                        infered = True
                        break
                if infered:
                    break
        if infered:
            return currStep
        else:
            return -1

    def move(self, feature):
        if self.randomLocation:
            locationOnObject = {
                "top": feature["top"] + np.random.rand() * feature["height"],
                "left": feature["left"] + np.random.rand() * feature["width"],
            }
        else:
            locationOnObject = {
                "top": feature["top"],
                "left": feature["left"]
            }
        if self.locationOnObject is not None:
            displacement = {"top": locationOnObject["top"] - self.locationOnObject["top"],
                            "left": locationOnObject["left"] - self.locationOnObject["left"]}
            if self.useNoise:
                self.corticalColumn.movementCompute(displacement, self.noiseFactor, self.moduleNoiseFactor)
            else:
                self.corticalColumn.movementCompute(displacement, 0, 0)
        self.locationOnObject = locationOnObject

    def sense(self, featureSDR):
        self.corticalColumn.sensoryCompute(featureSDR)

    def inferObject(self, objectDescription):
        self.reset()
        self.setLearnMode(False)
        currentStep = 0
        finished = False
        inferred = False
        inferredStep = None
        prevTouchSequence = None
        for _ in range(self.maxTraversals):
            while True:
                touchSequence = range(len(objectDescription["features"]))
                touchSequence = np.random.permutation(touchSequence)
                if prevTouchSequence is not None and touchSequence[0] == prevTouchSequence[-1]:
                    continue
                break
            for iFeature in touchSequence:
                currentStep += 1
                feature = objectDescription["features"][iFeature]
                self.move(feature)
                featureSDR = feature["featureValue"]
                self.sense(featureSDR)
                if not inferred:
                    representation = self.corticalColumn.getSensoryAssociatedLocationRepresentation()
                    target_representations = set(np.concatenate(self.locationRepresentations[(objectDescription["name"], iFeature)]))
                    inferred = (set(representation) == target_representations)
                    if inferred:
                        inferredStep = currentStep
                    if not inferred and tuple(representation) in self.representationSet:
                        print("Converged to an incorrect representation!")
                        return None
                finished = (inferred and self.numSensations is None) or (self.numSensations is not None and currentStep == self.numSensations)
                if finished:
                    break
            prevTouchSequence = touchSequence
            if finished:
                break
        return inferredStep

    def setLearnMode(self, isLearning):
        self.corticalColumn.setLearnMode(isLearning)

    def reset(self):
        self.corticalColumn.reset()
        self.locationOnObject = None
